import numpy as np
import pandas as pd
from scipy import signal
import time
from collections import deque

class AlertCalculator:
    """智能预警计算器，多维度融合评分模型"""
    def __init__(self, config=None):
        self.config = config or {
            'tech_weight': 0.6,      # 技术面权重
            'sentiment_weight': 0.25,  # 情绪面权重
            'liquidity_weight': 0.15,  # 流动性权重
            'alert_threshold': 0.7,   # 预警阈值
            'high_risk_threshold': 0.85,  # 高风险预警阈值
            'history_window': 100    # 历史记录窗口大小
        }
        
        # 历史预警记录
        self.alert_history = deque(maxlen=self.config['history_window'])
        self.score_history = deque(maxlen=self.config['history_window'])
        
        # 缓存最近的评分结果
        self.recent_scores = {}
        
        # 预警级别定义
        self.ALERT_LEVELS = {
            'NO_ALERT': {'level': 0, 'name': '无预警', 'color': 'green'},
            'LOW': {'level': 1, 'name': '低风险预警', 'color': 'blue'},
            'MEDIUM': {'level': 2, 'name': '中风险预警', 'color': 'yellow'},
            'HIGH': {'level': 3, 'name': '高风险预警', 'color': 'orange'},
            'CRITICAL': {'level': 4, 'name': '严重预警', 'color': 'red'}
        }
    
    def wavelet_analysis(self, prices, wavelet='db4', level=3):
        """使用小波分析进行价格波动特征提取"""
        if len(prices) < 32:  # 小波分析需要足够的数据点
            return 0.5  # 返回中性分数
        
        try:
            # 计算收益率
            returns = np.diff(prices) / prices[:-1]
            
            # 使用离散小波变换分解收益率序列
            # 由于scipy的dwt不直接支持多维分解，这里使用简化的波动分析
            # 在实际应用中，可以使用pywavelets库进行更复杂的小波分析
            
            # 计算波动率特征
            volatility = np.std(returns)
            
            # 计算趋势特征（使用简单移动平均）
            if len(prices) >= 10:
                short_ma = np.mean(prices[-5:])
                long_ma = np.mean(prices[-20:])
                trend_score = min(1.0, max(0.0, (short_ma / long_ma - 0.98) / 0.04))
            else:
                trend_score = 0.5
            
            # 计算价格变化加速度
            if len(returns) >= 5:
                acceleration = np.diff(returns[-5:])
                acceleration_score = min(1.0, max(0.0, (np.mean(abs(acceleration)) / (volatility * 0.1) + 1) / 2))
            else:
                acceleration_score = 0.5
            
            # 综合技术面评分
            tech_score = 0.4 * trend_score + 0.3 * (volatility / (np.mean(np.abs(returns)) + 1e-8)) + 0.3 * acceleration_score
            
            # 归一化到0-1范围
            tech_score = min(1.0, max(0.0, tech_score))
            
            return tech_score
        except Exception as e:
            print(f"Error in wavelet analysis: {str(e)}")
            return 0.5
    
    def calc_orderbook_imbalance(self, orderbook_data):
        """计算订单簿不平衡度"""
        if not orderbook_data or 'bids' not in orderbook_data or 'asks' not in orderbook_data:
            return 0.5  # 返回中性分数
        
        try:
            # 获取买单和卖单数据
            bids = orderbook_data['bids']
            asks = orderbook_data['asks']
            
            # 计算深度内的买单和卖单总量
            depth = min(5, len(bids), len(asks))  # 取前5档或更少
            
            bid_volume = sum(size for _, size in bids[:depth])
            ask_volume = sum(size for _, size in asks[:depth])
            
            total_volume = bid_volume + ask_volume
            
            if total_volume == 0:
                return 0.5
            
            # 计算不平衡度（归一化到0-1范围）
            imbalance = (bid_volume - ask_volume) / total_volume  # 范围[-1, 1]
            liquidity_score = (imbalance + 1) / 2  # 转换为[0, 1]
            
            # 考虑订单簿更新频率
            if 'update_frequency' in orderbook_data:
                # 更新频率越高，流动性越好
                frequency_score = min(1.0, orderbook_data['update_frequency'] / 100)  # 假设100次/秒为高频
                liquidity_score = 0.7 * liquidity_score + 0.3 * frequency_score
            
            return liquidity_score
        except Exception as e:
            print(f"Error calculating orderbook imbalance: {str(e)}")
            return 0.5
    
    def check_market_regime(self, regime_detector):
        """检查当前市场状态"""
        if regime_detector is None:
            return 'NEUTRAL'
        
        try:
            return regime_detector.current_regime
        except:
            return 'NEUTRAL'
    
    def calculate_alert_level(self, stock_data, regime_detector=None):
        """计算预警级别"""
        if not stock_data or 'prices' not in stock_data:
            return {
                'alert_level': 'NO_ALERT',
                'alert_score': 0.0,
                'reason': 'Insufficient data'
            }
        
        # 记录计算开始时间
        start_time = time.time()
        
        # 1. 计算技术面评分
        tech_score = self.wavelet_analysis(stock_data['prices'])
        
        # 2. 计算情绪面评分
        if 'news' in stock_data and stock_data['news']:
            # 简单的情绪评分计算，实际应用中可以使用BERT等模型
            sentiment_scores = []
            for news_item in stock_data['news']:
                if 'sentiment_score' in news_item:
                    # 假设sentiment_score范围是[-1, 1]
                    normalized_score = (news_item['sentiment_score'] + 1) / 2  # 转换为[0, 1]
                    sentiment_scores.append(normalized_score)
                elif 'sentiment' in news_item:
                    # 映射情感标签到分数
                    sentiment_map = {'positive': 1.0, 'neutral': 0.5, 'negative': 0.0}
                    sentiment_scores.append(sentiment_map.get(news_item['sentiment'], 0.5))
            
            if sentiment_scores:
                sentiment_score = np.mean(sentiment_scores)
            else:
                sentiment_score = 0.5
        else:
            sentiment_score = 0.5
        
        # 3. 计算流动性评分
        if 'orderbook' in stock_data:
            liquidity_score = self.calc_orderbook_imbalance(stock_data['orderbook'])
        else:
            liquidity_score = 0.5
        
        # 4. 融合多维度评分
        alert_score = (
            self.config['tech_weight'] * tech_score +
            self.config['sentiment_weight'] * sentiment_score +
            self.config['liquidity_weight'] * liquidity_score
        )
        
        # 5. 检查市场状态，极端波动行情下可能需要调整预警阈值
        market_regime = self.check_market_regime(regime_detector)
        
        # 确定预警级别
        if market_regime == 'EXTREME_VOLATILITY':
            # 极端波动行情下，降低预警阈值以提高敏感度
            adjusted_threshold = self.config['alert_threshold'] * 0.8
        else:
            adjusted_threshold = self.config['alert_threshold']
        
        if alert_score >= self.config['high_risk_threshold']:
            alert_level = 'CRITICAL'
        elif alert_score >= self.config['alert_threshold']:
            alert_level = 'HIGH'
        elif alert_score >= adjusted_threshold * 0.8:
            alert_level = 'MEDIUM'
        elif alert_score >= adjusted_threshold * 0.6:
            alert_level = 'LOW'
        else:
            alert_level = 'NO_ALERT'
        
        # 生成预警原因
        reasons = []
        if tech_score > 0.7:
            reasons.append(f'技术面指标异常 (分数: {tech_score:.2f})')
        if sentiment_score > 0.8 or sentiment_score < 0.2:
            reasons.append(f'市场情绪极端 (分数: {sentiment_score:.2f})')
        if liquidity_score > 0.8 or liquidity_score < 0.2:
            reasons.append(f'流动性失衡 (分数: {liquidity_score:.2f})')
        if market_regime == 'EXTREME_VOLATILITY':
            reasons.append('当前市场处于极端波动状态')
        
        # 记录评分历史
        current_time = int(time.time() * 1000)
        symbol = stock_data.get('symbol', 'unknown')
        
        self.score_history.append({
            'timestamp': current_time,
            'symbol': symbol,
            'tech_score': tech_score,
            'sentiment_score': sentiment_score,
            'liquidity_score': liquidity_score,
            'alert_score': alert_score
        })
        
        # 缓存最近评分结果
        self.recent_scores[symbol] = {
            'tech_score': tech_score,
            'sentiment_score': sentiment_score,
            'liquidity_score': liquidity_score,
            'timestamp': current_time
        }
        
        # 如果达到预警级别，记录预警历史
        if alert_level != 'NO_ALERT':
            alert_record = {
                'timestamp': current_time,
                'symbol': symbol,
                'alert_level': alert_level,
                'alert_score': alert_score,
                'tech_score': tech_score,
                'sentiment_score': sentiment_score,
                'liquidity_score': liquidity_score,
                'market_regime': market_regime,
                'reasons': reasons,
                'calculation_time_ms': (time.time() - start_time) * 1000
            }
            
            self.alert_history.append(alert_record)
            
            # 触发预警（实际应用中可以连接到通知系统）
            if alert_level in ['HIGH', 'CRITICAL']:
                self._trigger_alert(alert_record)
            
            return alert_record
        
        return {
            'alert_level': alert_level,
            'alert_score': alert_score,
            'tech_score': tech_score,
            'sentiment_score': sentiment_score,
            'liquidity_score': liquidity_score,
            'market_regime': market_regime,
            'reasons': reasons,
            'timestamp': current_time,
            'calculation_time_ms': (time.time() - start_time) * 1000
        }
    
    def _trigger_alert(self, alert_record):
        """触发预警通知"""
        # 实际应用中，这里可以连接到消息队列、邮件系统或短信服务
        print(f"ALERT: {alert_record['symbol']} - {self.ALERT_LEVELS[alert_record['alert_level']]['name']} - Score: {alert_record['alert_score']:.2f}")
        print(f"Reasons: {', '.join(alert_record['reasons'])}")
    
    def batch_calculate_alerts(self, stocks_data, regime_detector=None):
        """批量计算多个股票的预警级别"""
        results = {}
        alerts_triggered = 0
        
        start_time = time.time()
        
        for symbol, stock_data in stocks_data.items():
            # 确保股票数据包含符号信息
            stock_data_with_symbol = stock_data.copy()
            stock_data_with_symbol['symbol'] = symbol
            
            # 计算预警级别
            result = self.calculate_alert_level(stock_data_with_symbol, regime_detector)
            results[symbol] = result
            
            # 统计触发的预警数量
            if result.get('alert_level', 'NO_ALERT') != 'NO_ALERT':
                alerts_triggered += 1
        
        total_time = (time.time() - start_time) * 1000  # 毫秒
        avg_time_per_stock = total_time / len(stocks_data) if stocks_data else 0
        
        return {
            'results': results,
            'total_stocks': len(stocks_data),
            'alerts_triggered': alerts_triggered,
            'total_calculation_time_ms': total_time,
            'avg_calculation_time_per_stock_ms': avg_time_per_stock,
            'timestamp': int(time.time() * 1000)
        }
    
    def get_alert_statistics(self, lookback=None):
        """获取预警统计信息"""
        if not self.alert_history:
            return {
                'total_alerts': 0,
                'alert_distribution': {},
                'most_alerted_symbols': []
            }
        
        # 获取指定时间范围内的预警记录
        if lookback:
            current_time = int(time.time() * 1000)
            lookback_time = current_time - (lookback * 60 * 1000)  # 转换为毫秒
            filtered_history = [alert for alert in self.alert_history if alert['timestamp'] >= lookback_time]
        else:
            filtered_history = list(self.alert_history)
        
        # 计算预警分布
        alert_counts = {level: 0 for level in self.ALERT_LEVELS if level != 'NO_ALERT'}
        symbol_counts = {}
        
        for alert in filtered_history:
            level = alert['alert_level']
            if level in alert_counts:
                alert_counts[level] += 1
            
            symbol = alert['symbol']
            symbol_counts[symbol] = symbol_counts.get(symbol, 0) + 1
        
        # 找出预警最多的股票
        most_alerted = sorted(symbol_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        
        # 计算平均预警分数
        avg_alert_scores = {}
        score_sums = {level: 0 for level in alert_counts}
        
        for alert in filtered_history:
            level = alert['alert_level']
            score_sums[level] += alert['alert_score']
        
        for level, count in alert_counts.items():
            if count > 0:
                avg_alert_scores[level] = score_sums[level] / count
            else:
                avg_alert_scores[level] = 0
        
        return {
            'total_alerts': len(filtered_history),
            'alert_distribution': {
                level: {
                    'count': count,
                    'name': self.ALERT_LEVELS[level]['name'],
                    'avg_score': avg_alert_scores[level]
                }
                for level, count in alert_counts.items()
            },
            'most_alerted_symbols': most_alerted,
            'time_range': lookback if lookback else 'all_time',
            'timestamp': int(time.time() * 1000)
        }
    
    def get_recent_scores(self, symbol):
        """获取指定股票的最近评分结果"""
        return self.recent_scores.get(symbol, None)
    
    def set_weights(self, tech_weight=None, sentiment_weight=None, liquidity_weight=None):
        """动态调整各维度的权重"""
        if tech_weight is not None:
            self.config['tech_weight'] = tech_weight
        if sentiment_weight is not None:
            self.config['sentiment_weight'] = sentiment_weight
        if liquidity_weight is not None:
            self.config['liquidity_weight'] = liquidity_weight
        
        # 确保权重和为1
        total_weight = (self.config['tech_weight'] + 
                       self.config['sentiment_weight'] + 
                       self.config['liquidity_weight'])
        
        if total_weight != 1.0:
            # 归一化权重
            self.config['tech_weight'] /= total_weight
            self.config['sentiment_weight'] /= total_weight
            self.config['liquidity_weight'] /= total_weight
        
        return {
            'tech_weight': self.config['tech_weight'],
            'sentiment_weight': self.config['sentiment_weight'],
            'liquidity_weight': self.config['liquidity_weight']
        }